#!/usr/bin/env python
from __future__ import print_function
import numpy as np
from astropy.io import fits as pf

class SimpleEBoundObj(object):

    def __init__(self, infile):
        if not isinstance(infile, str):
            raise IOError('infile should be a string!')
        self.infile = infile
        hdulist = pf.open(self.infile, memmap=True)

        try:
            hdu = hdulist['EBOUNDS']
        except KeyError:
            try:
                hdu = hdulist['ENERGIES']
            except KeyError:
                try:
                    hdu = hdulist['ENERGYBINS']
                except KeyError:
                    raise IOError('No valid extension (EBOUNDS|ENERGIES|ENERGYBINS)!')
                else:
                    self.ebType = 'ccube'
            else:
                self.ebType = 'expcube'
        else:
            self.ebType = 'ccube'

        tunit2scale = {'eV': 1e-6, 'keV': 1e-3, 'MeV': 1., 'GeV': 1e3, 'TeV': 1e6}
        if self.ebType == 'ccube':
            emin, emax = None, None
            for iname, name in enumerate(hdu.data.names):
                if name.lower() == 'e_min':
                    try:
                        tunit = hdu.header['TUNIT%i'%(iname+1)]
                    except KeyError:
                        scale = 1.
                    else:
                        if tunit in tunit2scale:
                            scale = tunit2scale[tunit]
                        else:
                            scale = 1.
                    emin = np.array(hdu.data.field(iname) * scale).flatten() # [MeV]

                    if emax is not None:
                        break
                elif name.lower() == 'e_max':
                    try:
                        tunit = hdu.header['TUNIT%i'%(iname+1)]
                    except KeyError:
                        scale = 1.
                    else:
                        if tunit in tunit2scale:
                            scale = tunit2scale[tunit]
                        else:
                            scale = 1.
                    emax = np.array(hdu.data.field(iname) * scale).flatten() # [MeV]

                    if emin is not None:
                        break

            if (emin is None) or (emax is None):
                raise RuntimeError('No "E_MIN" or "E_MAX" in the EBOUNDS extension!')
            self.Energy = np.unique(np.r_[emin, emax]) # [MeV]
        elif self.ebType == 'expcube':
            self.Energy = None
            for iname, name in enumerate(hdu.data.names):
                if name.lower() == 'energy':
                    try:
                        tunit = hdu.header['TUNIT%i'%(iname+1)]
                    except KeyError:
                        scale = 1.
                    else:
                        if tunit in tunit2scale:
                            scale = tunit2scale[tunit]
                        else:
                            scale = 1.
                    self.Energy = hdu.data.field(iname)*scale # [MeV]

                    break

            if self.Energy is None:
                raise RuntimeError('No "Energy" in the ENERGIES extension!')

        del hdu.data

    @property
    def Emin(self):
        return self.Energy[:-1]

    @property
    def Emax(self):
        return self.Energy[1:]

    @property
    def nEBound(self):
        return self.Energy.size-1 # No. of Ebins

    @property
    def nEnergy(self):
        return self.Energy.size # No. of energy boundary

def cli():
    import argparse

    parser = argparse.ArgumentParser(description='Show ebins!')
    parser.add_argument("infiles", type=str, nargs='+', help='input file(s)')
    args = parser.parse_args()

    for infile in args.infiles:
        try:
            eb = SimpleEBoundObj(infile)
        except IOError:
            print('No energy info in %s!'%infile)
        else:
            if eb.ebType == 'ccube':
                print('Energy bins in %s:' % infile)
                for ie, _emin, _emax in zip(range(eb.nEBound), eb.Emin, eb.Emax):
                    print('BIN%02i: (%r, %r) MeV' % (ie, _emin, _emax))
            else:
                print('Energy layers in %s:' % infile)
                for ie, e in enumerate(eb.Energy):
                    print('LAYER%02i: %r MeV' % (ie, e))

if __name__ == '__main__':
    cli()
